/*
 * Matrix.cpp
 *
 *  Created on: Apr 9, 2010
 *      Author: andrew
 */

#include <iostream>
#include "Matrix.h"
using namespace std;

/**
 * Constructor
 */
Matrix::Matrix(int rows, int cols) {
	m = rows;
	n = cols;
	matrix = new double[rows*cols];
}

/*
 * Default constructor
 */
Matrix::Matrix() {
}

/**
 * Destructor
 */
Matrix::~Matrix() {
	delete[] matrix;
}

/**
 * Copy constructor
 */
Matrix::Matrix(const Matrix& copyme){
	n = copyme.n;
	m = copyme.m;
	matrix = new double[n*m];
	for(int i = 0; i < n*m; i++){
		matrix[i] = copyme.matrix[i];
	}
}

/**
 * Assignment operator
 */
Matrix& Matrix::operator=(const Matrix &mObj){
    if (this != &mObj)
    {
        delete[] matrix;
        m = mObj.m;
        n = mObj.n;
        matrix = new double[n*m];
        for(int i = 0; i < n*m; i++){
			matrix[i] = mObj.matrix[i];
		}
    }
    return *this;
}

/**
 * index entries in matrix
 */
double& Matrix::operator()(int i, int j) const{
	return matrix[i*n + j];
}

/**
 * Vector operator
 * i = first row index, j = final row index
 * k = column index
 */
Vector1D Matrix::vecNot(int idx1, int idx2, int k) const{

	Vector1D x(idx2-idx1+1);
	x.initConst(0.0);

	if ((idx1 > m) || (idx2 > m)) {
		cout << "Row index exceeds number of rows in the matrix" << endl;
		return x;
	}
	if (k > n) {
		cout << "Column index exceeds number of columns in the matrix" << endl;
		return x;
	}

	int ctr = 0;
	for(int i = idx1; i <= idx2; i++) {
		x(ctr) = (*this)(i,k); ctr++;
	}
	return x;
}

/*
 * Returns exp(t*A)
Matrix Matrix::expm(double tstep) const {

	if(n!=m) {
		cout << "Matrix must be square" << endl;
		return *this;
	}

	double* A_reorder = c2fortran(*this);

	double *wsp, *mVec;
	int ideg, lwsp, iexph, ns, info;
	int *ipiv;

	ideg = 6;   // The degree of the diagonal Pade to be used (6 is usually good enough)
	ns = 0;     // # of scaling-squaring used
	lwsp = 4*m*m+ideg+1; // Workspace length;

	wsp  = new double[lwsp];
	ipiv = new int[m];
	mVec = new double[lwsp];

	dgpadm_(&ideg,&m,&tstep,A_reorder,&m,wsp,&lwsp,ipiv,&iexph,&ns,&info);

	// Extract solution from array wsp, solution starts at iexph
	for (int i = 0; i < m*m; i++) {
		*(mVec+i) = *(wsp+iexph-1+i);
	}

	Matrix expmA = fortran2c(mVec,m,n);

	delete[] wsp; delete[] ipiv; delete[] mVec; delete[] A_reorder;

	return expmA;
}
*/

/*
 * Appends a matrix at a specified row
 */
Matrix Matrix::appendMatrixBelow(Matrix& A) const{
	if (A.n!=n) {
		cout << "Matrices must have the same number of columns" << endl;
		return *this;
	}

	Matrix Anew = Matrix(m+A.m,n); int k = 0;
	for (int i = 0; i < m; i++) {
		for (int j = 0; j < n; j++) {
			Anew(i,j) = matrix[k]; k++;
		}
	}
	for (int i = 0; i < A.m; i++) {
		for (int j = 0; j < n; j++) {
			Anew(i+m,j) = A(i,j);
		}
	}

	return Anew;
}

/*
 * append a matrix to the side
 */
Matrix Matrix::appendMatrixSide(Matrix& A) const{
	if (A.m!=m) {
		cout << "Matrices must have the same number of rows" << endl;
		return *this;
	}

	Matrix Anew = Matrix(m,n+A.n);
	for (int i = 0; i < m; i++) {
		for (int j = 0; j < n; j++) {
			Anew(i,j) = operator ()(i,j);
		}
	}
	for (int i = 0; i < m; i++) {
		for (int j = 0; j < A.n; j++) {
			Anew(i,j+n) = A(i,j);
		}
	}

	return Anew;
}

/*
 * Returns a column from a matrix
 */
Vector1D Matrix::selectCol(int col) const {
	Vector1D x(m);
	for(int i = 0; i < m; i++) {
		x(i) = (*this)(i,col);
	}
	return x;
}

/*
 * Creates a diagonal matrix
 */
Matrix Matrix::diag(Vector1D d) const {
	Matrix D(d.length(),d.length());
	D.initConst(0.0);
	for(int i = 0; i < d.length(); i++) {
		D(i,i) = d(i);
	}
	return D;
}

/**
 * append a vector at a specified column
 */
Matrix Matrix::appendVectorCol(Vector1D& v, int col) const{
//	cout << m << endl;
//	cout << n << endl;
//	cout << v.length() << endl;
	if (v.length()!=m) {
		cout << "Matrix must have the same number of rows as vector" << endl;
		return *this;
	}

	for (int i = 0; i < m; i++) {
		(*this)(i,col) = v(i);
	}

	return *this;
}

/*
 * Subselects a smaller matrix from the current one. Inputs 1 & 2 indicate the starting row and column
 * Inputs 3 & 4 indicate how many rows and columns to select
 */
Matrix Matrix::subselect(int rowStart, int colStart, int nrows, int ncols) const {
	if ((rowStart+nrows)>m || (colStart+ncols)>n) {
		cout << "Out of bounds error in subselection" << endl;
		return *this;
	}
	Matrix S(nrows,ncols);
	for (int i = rowStart; i < nrows+rowStart; i++)
		for (int j = colStart; j < ncols+colStart; j++)
			S(i-rowStart,j-colStart) = operator()(i,j);

	return S;
}

/*
 * Initialize a matrix to a constant value
 */
Matrix Matrix::initConst(double init) const{

	for(int i = 0; i < m*n; i++) matrix[i] = init;

	return *this;
}

/**
 * return transpose of matrix
 */
Matrix Matrix::getTranspose() const{

	Matrix AT = Matrix(n,m);
	for(int i = 0; i < m; i++)
		for(int j = 0; j < n; j++)
			AT(j,i) = operator ()(i,j);
	return AT;
}

/*
 * Matrix matrix addition
 */
Matrix Matrix::operator+(const Matrix& B) const {
	if (B.m!=m || B.n!=n) {
		cout << "Matrix sizes are incompatible for addition" << endl;
		return *this;
	}
	Matrix ApB(m,n);
	for (int i = 0; i < m; i++)
		for (int j = 0; j < n; j++)
			ApB(i,j) = operator()(i,j) + B(i,j);

	return ApB;
}

/*
 * Matrix matrix subtraction
 */
Matrix Matrix::operator-(const Matrix& B) const {
	if (B.m!=m || B.n!=n) {
		cout << "Matrix sizes are incompatible for subtraction" << endl;
		return *this;
	}
	Matrix AmB(m,n);
	for (int i = 0; i < m; i++)
		for (int j = 0; j < n; j++)
			AmB(i,j) = operator()(i,j) - B(i,j);

	return AmB;
}

/**
 * matrix-matrix multiply
 */
Matrix Matrix::operator*(const Matrix& B) const{

	if (B.m!=n) {
		cout << "Matrix sizes are incompatible for multiplication" << endl;
		return *this;
	}

	// reorder matrices fortran style
	double* A_reorder = c2fortran(*this);
	double* B_reorder = c2fortran(B);

	//call blas
	double done = 1.0;
	double dzero = 0.0;
	double *Cb = new double[m*B.n];
	dgemm_("n","n",&m,&B.n,&n,&done,A_reorder,&m,B_reorder,&B.m,&dzero,Cb,&m);

	// Reorder vector from FORTRAN form
	Matrix C = fortran2c(Cb,m,B.n);

	delete[] Cb; delete[] A_reorder; delete[] B_reorder;

	return C;
}


/**
 * matrix-vector multiply
 */
Vector1D Matrix::operator*(const Vector1D& x) const{
	double* y = new double[m];
	double* b = x.getCopyOfArray();

	// unroll matrix in fortran "column-major" style
	double* Areorder = c2fortran(*this);

	// call blas
	int ione = 1;
	double done = 1.0;
	double dzero = 0.0;
	dgemv_("n",&m,&n,&done,Areorder,&m,b,&ione,&dzero,y,&ione);

	delete[] Areorder;
	delete[] b;

	// assign array to vector object
	Vector1D yvec = Vector1D(m,y);
	delete[] y;

	return yvec;
}

/*
 * Matrix scalar multiply
 */
Matrix Matrix::operator *(double s) const {
	Matrix As(m,n);
	for (int i = 0; i < m*n; i++){
		As.matrix[i] = matrix[i]*s;
	}

	return As;
}

/*
 * Places matrix B below and to the right
 */
Matrix Matrix::blkdiag(const Matrix &B) const {
	Matrix Anew(m+B.m,n+B.n);
	Anew = Anew.initConst(0);

	for(int i = 0; i < m; i++)
		for(int j = 0; j < n; j++)
			Anew(i,j) = operator()(i,j);
	for(int i = m; i < m+B.m; i++)
		for(int j = n; j < n+B.n; j++)
			Anew(i,j) = B(i-m,j-n);

	return Anew;
}

/**
 * solve linear system of equations
 */
Vector1D Matrix::solve(const Vector1D& b) const{

	double* b_array = b.getCopyOfArray();

	// unroll matrix in fortran "column-major" style
	double* Areorder = c2fortran(*this);

	// call lapack LU decomposition
	int ione = 1;
	int ipiv[m];
	int info;
	dgesv_(&m,&ione,Areorder,&m,ipiv,b_array,&m,&info);

	delete[] Areorder;

	// assign array to a vector object
	Vector1D yvec = Vector1D(m,b_array);
	delete[] b_array;

	return yvec;
}

/*
 * Compute matrix inverse
 */
Matrix Matrix::inv() const{

	if (m != n) {
		cout << "Matrix must be square to invert" << endl;
		return *this;
	}

	// unroll matrix in fortran "column-major" style
	double* Areorder = c2fortran(*this);

    int *IPIV = new int[m];
    int LWORK = m*m;
    double *WORK = new double[LWORK];
    int INFO;
    int N = m;

    dgetrf_(&N,&N,Areorder,&N,IPIV,&INFO);
    dgetri_(&N,Areorder,&N,IPIV,WORK,&LWORK,&INFO);

	// Reorder vector from FORTRAN form
	Matrix B = fortran2c(Areorder,m,m);

	delete[] Areorder;

	return B;
}

/**
 * solve linear system of equations with a matrix input
 */
Matrix Matrix::solve(const Matrix& B) const{

	// reorder matrices fortran style
	double* A_reorder = c2fortran(*this);
	double* B_reorder = c2fortran(B);

	// call lapack LU decomposition
	int ipiv[m];
	int info;
	dgesv_(&m,&B.n,A_reorder,&m,ipiv,B_reorder,&B.m,&info);
	if (info != 0) {
		cout << "Error - LU decomposition failed with error = " << info << endl;
	}

	// Reorder vector from FORTRAN form
	Matrix C = fortran2c(B_reorder,B.m,B.n);

	delete[] A_reorder; delete[] B_reorder;

	return C;
}

Matrix Matrix::eye(int n){

	Matrix A(n,n);
	A.initConst(0.0);

	for (int i = 0; i < n; i++){
		A(i,i) = 1.0;
	}
	return A;

}

Matrix Matrix::fortran2c(const double *B, int n, int m){
	Matrix A = Matrix(n,m);
	int k = 0;
	for (int j = 0; j < m; j++) {
		for(int i = 0; i < n; i++) {
			A(i,j) = *(B+k); k++;
		}
	}
	return A;
}

Matrix Matrix::chol() const {

	// reorder matrix fortran style
	double* A_reorder = c2fortran(*this);

	// call lapack LU decomposition
	int info;
	dpotrf_("U",&m,A_reorder,&m,&info);
	if (info != 0) {
		cout << "Error - Cholesky factorization failed with error = " << info << endl;
	}

	// Reorder vector from FORTRAN form
	Matrix Chol = fortran2c(A_reorder,m,n);

	delete[] A_reorder;

	return Chol;
}

double* Matrix::c2fortran(const Matrix A){
	double* Areorder = new double[A.n*A.m];

	// unroll matrix in fortran "column-major" style
	for (int i = 0; i < A.m; i++)
		for(int j = 0; j < A.n; j++)
			Areorder[A.m*j+i] = A(i,j);

	return Areorder;
}

Matrix Matrix::selectEntries(const Vector1D& selectRows, const Vector1D &selectCols){
	int nrows = 0; int ncols = 0;
	for (int i = 0; i < selectRows.length(); i++) nrows += selectRows(i);
	for (int i = 0; i < selectCols.length(); i++) ncols += selectCols(i);
	Matrix A = Matrix(nrows,ncols);
	int m = 0; int n = 0;
	for(int i = 0; i < selectRows.length(); i++){
		for(int j = 0; j < selectCols.length(); j++){
			if (selectRows(i) && selectCols(j)){
				A(m,n) = (*this)(i,j);
				n++;
			}
		}
		if (selectRows(i)) m++;
		n = 0;
	}
	return A;
}

void Matrix::printMatrix(void) const {
	for (int i = 0; i < m; i++) {
		for (int j = 0; j < n; j++) {
			cout << (*this)(i,j) << " ";
		}
		cout << endl;
	}
}

void Matrix::printMatrix(int rowstart, int rowend, int colstart, int colend) const {
	if ((rowstart > m) || (rowend > m)) {
		cout << "Rows specified in printMatrix exceeds number in Matrix" << endl;
	}
	if ((colstart > n) || (colend > n)) {
		cout << "Columns specified in printMatrix exceeds number in Matrix" << endl;
	}
	if ((rowstart < 0) || (rowend < 0)) {
		cout << "Rows specified in printMatrix must exceed zero" << endl;
	}
	if ((colstart < 0) || (colend < 0)) {
		cout << "Columns specified in printMatrix must exceed zero" << endl;
	}
	for (int i = rowstart; i <= rowend; i++) {
		for (int j = colstart; j <= colend; j++) {
			cout << (*this)(i,j) << " ";
		}
		cout << endl;
	}
}
